[mlir][linalg] Allow promotion to use the original subview size #144334
Conversation
zbenzion
requested review from
ftynse,
nicolasvasilache,
dcaballe and
rengolin
as code owners
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-linalg Author: None (zbenzion) Changeslinalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic. Adding a promotion option to use the original subview size by default, thus minimizing the allocation size. Full diff: https://github.com/llvm/llvm-project/pull/144334.diff 5 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 15ea5e7bf7159..42de4a2f153c9 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1309,6 +1309,7 @@ def PromoteOp : Op<Transform_Dialect, "structured.promote",
DefaultValuedAttr<I64ArrayAttr, "{}">:$operands_to_promote,
DefaultValuedAttr<BoolArrayAttr, "{}">:$use_full_tile_buffers,
UnitAttr:$use_full_tiles_by_default,
+ UnitAttr:$use_original_subview_size,
UnitAttr:$use_alloca,
OptionalAttr<AnyAttr>:$memory_space,
OptionalAttr<DeviceMappingArrayAttr>:$mapping,
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 2eef0a06d0eb4..d9dceade0ac95 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -381,6 +381,13 @@ struct LinalgPromotionOptions {
useFullTileBuffersDefault = use;
return *this;
}
+ /// If true bufffers will be allocated with the original subview size. This
+ /// may result in more dynamic allocations, in case of dynamic sizes.
+ bool useOriginalSubviewSize = false;
+ LinalgPromotionOptions &setUseOriginalSubviewSize(bool originalSize) {
+ useOriginalSubviewSize = originalSize;
+ return *this;
+ }
/// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
std::optional<unsigned> alignment;
LinalgPromotionOptions &setAlignment(unsigned align) {
@@ -712,7 +719,8 @@ FailureOr<LinalgOp> specializeGenericOp(RewriterBase &rewriter,
GenericOp genericOp);
/// Create a new buffer using the `allocationFn` provided. The size of this
-/// buffer is the smallest constant bounding size along each dimension that
+/// buffer is either the original subview size when 'useOriginalSubviewSize' is
+/// set to true or the smallest constant bounding size along each dimension that
/// can be computed for the size of the result of `subView`. Returns the
/// allocated buffer as `fullLocalView` and the view that matches the size of
/// the result of subview operation as `partialLocalView`.
@@ -722,6 +730,7 @@ struct PromotionInfo {
};
FailureOr<PromotionInfo>
promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
+ bool useOriginalSubviewSize,
const AllocBufferCallbackFn &allocationFn,
DataLayout &layout);
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index b2c28f5eed33c..6be73a11e1b54 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2236,6 +2236,9 @@ transform::PromoteOp::applyToOne(transform::TransformRewriter &rewriter,
if (getUseFullTilesByDefault())
promotionOptions = promotionOptions.setUseFullTileBuffersByDefault(
getUseFullTilesByDefault());
+ if (getUseOriginalSubviewSize())
+ promotionOptions =
+ promotionOptions.setUseOriginalSubviewSize(getUseOriginalSubviewSize());
if (getUseAlloca())
promotionOptions = promotionOptions.setUseAlloca(getUseAlloca());
if (!getUseFullTileBuffers().empty())
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index d599ddd220dde..80b37544a99f8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -148,6 +148,9 @@ struct LinalgOpInstancePromotionOptions {
llvm::SmallSet<int64_t, 4> operandsNumbersToCopyIn;
/// True if the full view should be used for the promoted buffer.
DenseMap<Value, bool> useFullTileBuffers;
+ /// True if the original subview size should be used. This means the full tile
+ /// buffer is the same size as the partial view.
+ bool useOriginalSubviewSize;
/// Callback functions for allocation and deallocation of promoted buffers, as
/// well as to copy the data into and out of these buffers.
@@ -170,6 +173,7 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
options.useFullTileBuffers.value_or(llvm::SmallBitVector());
vUseFullTileBuffers.resize(linalgOp->getNumOperands(),
options.useFullTileBuffersDefault);
+ useOriginalSubviewSize = options.useOriginalSubviewSize;
for (OpOperand &opOperand : linalgOp->getOpOperands()) {
int64_t operandNumber = opOperand.getOperandNumber();
@@ -237,7 +241,8 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
// by a partial `copy` op.
FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
OpBuilder &b, Location loc, memref::SubViewOp subView,
- const AllocBufferCallbackFn &allocationFn, DataLayout &layout) {
+ bool useOriginalSubviewSize, const AllocBufferCallbackFn &allocationFn,
+ DataLayout &layout) {
auto viewType = subView.getType();
auto rank = viewType.getRank();
SmallVector<Value, 4> fullSizes;
@@ -254,7 +259,8 @@ FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
// to look for the bound.
LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
Value size;
- if (auto attr = llvm::dyn_cast_if_present<Attribute>(rangeValue.size)) {
+ if (llvm::isa_and_present<Attribute>(rangeValue.size) ||
+ useOriginalSubviewSize) {
size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
} else {
FailureOr<int64_t> upperBound =
@@ -295,7 +301,8 @@ promoteSubViews(ImplicitLocOpBuilder &b,
memref::SubViewOp subView =
cast<memref::SubViewOp>(v.second.getDefiningOp());
auto promotionInfo = promoteSubviewAsNewBuffer(
- b, b.getLoc(), subView, options.allocationFn, layout);
+ b, b.getLoc(), subView, options.useOriginalSubviewSize,
+ options.allocationFn, layout);
if (failed(promotionInfo))
return failure();
promotionInfoMap[v.first] = *promotionInfo;
diff --git a/mlir/test/Dialect/Linalg/promotion_options.mlir b/mlir/test/Dialect/Linalg/promotion_options.mlir
index caa72ba24316f..dbc073c2665f9 100644
--- a/mlir/test/Dialect/Linalg/promotion_options.mlir
+++ b/mlir/test/Dialect/Linalg/promotion_options.mlir
@@ -42,3 +42,59 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+func.func @matmul_f32(%A: memref<512x256xf32>, %B: memref<256x512xf32>, %C: memref<256x256xf32>, %s0: index, %s1: index, %s2: index) {
+ %c0 = arith.constant 0 : index
+ %c256 = arith.constant 256 : index
+ %c512 = arith.constant 512 : index
+ scf.for %arg4 = %c0 to %c512 step %s0 {
+ scf.for %arg5 = %c0 to %c512 step %s1 {
+ scf.for %arg6 = %c0 to %c256 step %s2 {
+ %i0 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg4)[%s0]
+ %i1 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg5)[%s1]
+ %i2 = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%arg6)[%s2]
+ %0 = memref.subview %A[%arg4, %arg6][%i0, %i2][1, 1] : memref<512x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>
+ %1 = memref.subview %B[%arg6, %arg5][%i2, %i1][1, 1] : memref<256x512xf32> to memref<?x?xf32, strided<[512, 1], offset: ?>>
+ %2 = memref.subview %C[%arg4, %arg5][%i0, %i1][1, 1] : memref<256x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>
+ linalg.matmul
+ ins(%0, %1: memref<?x?xf32, strided<[256, 1], offset: ?>>,
+ memref<?x?xf32, strided<[512, 1], offset: ?>>)
+ outs(%2: memref<?x?xf32, strided<[256, 1], offset: ?>>)
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_f32(
+// CHECK-SAME: %[[ARG0:.*]]: memref<512x256xf32>
+// CHECK-SAME: %[[ARG1:.*]]: memref<256x512xf32>
+// CHECK-SAME: %[[ARG2:.*]]: memref<256x256xf32>
+// CHECK-SAME: %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+
+// CHECK: %[[i0:.*]] = affine.min
+// CHECK: %[[i1:.*]] = affine.min
+// CHECK: %[[i2:.*]] = affine.min
+
+// CHECK: %[[VAL_13:.*]] = arith.muli %[[i0]], %[[i2]] : index
+// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_13]], %[[C4]] : index
+// CHECK: %[[VAL_15:.*]] = memref.alloc(%[[VAL_14]]) : memref<?xi8>
+
+// CHECK: %[[VAL_18:.*]] = arith.muli %[[i2]], %[[i1]] : index
+// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_18]], %[[C4]] : index
+// CHECK: %[[VAL_20:.*]] = memref.alloc(%[[VAL_19]]) : memref<?xi8>
+
+// CHECK: %[[VAL_23:.*]] = arith.muli %[[i0]], %[[i1]] : index
+// CHECK: %[[VAL_24:.*]] = arith.muli %[[VAL_23]], %[[C4]] : index
+// CHECK: %[[VAL_25:.*]] = memref.alloc(%[[VAL_24]]) : memref<?xi8>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.promote %0 { use_original_subview_size } : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
|
zbenzion
force-pushed
the
linalg/promote-orig-size
branch
from
e990a31 to
629ff4a
Compare
| @@ -427,6 +427,13 @@ struct LinalgPromotionOptions { | |||
| useFullTileBuffersDefault = use; | |||
| return *this; | |||
| } | |||
| /// If true bufffers will be allocated with the original subview size. This | |||
There was a problem hiding this comment.
| /// If true bufffers will be allocated with the original subview size. This | |
| /// If true, buffers will be allocated with the original subview size. This |
| @@ -811,6 +819,7 @@ struct PromotionInfo { | |||
| }; | |||
| FailureOr<PromotionInfo> | |||
| promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView, | |||
| bool useOriginalSubviewSize, | |||
There was a problem hiding this comment.
Let's have this take const LinalgPromotionOptions & instead of an individual boolean.
There was a problem hiding this comment.
In general, sounds good!
It seems though that the caller promoteSubViews is already using LinalgOpInstancePromotionOptions, which is an internal representation in Promotion.cpp.
Another issue is that LinalgPromotionOptions doesn't force an explicit allocationFn, and promoteSubviewAsNewBuffer does.
We could maybe, alternatively, provide a LinalgPromoteSubviewAsNewBufferOptions, which will pack the boolean and the allocation function.
What do you think?
There was a problem hiding this comment.
Can we just extend one or another? I'd rather avoid yet another helper class so if it's impossible to extend one of the existing classes, we can stick with the single argument here.
There was a problem hiding this comment.
Since this method is used both internally in Promotion.cpp (LinalgOpInstancePromotionOptions) and is exposed externally (LinalgPromotionOptions) I feel like it'd be a little unnatural to use one or the other at this stage. I kept the single argument, as you suggested.
If everything looks good, I'd appreciate if you can hit the merge button for me, since I don't have write access.
Thanks for the review!
linalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic. Adding a promotion option to use the original subview size by default, thus minimizing the allocation size.
zbenzion
force-pushed
the
linalg/promote-orig-size
branch
from
629ff4a to
981ee07
Compare
|
@zbenzion Congratulations on having your first Pull Request (PR) merged into the LLVM Project! Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR. Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues. How to do this, and the rest of the post-merge process, is covered in detail here. If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again. If you don't get any reports, no action is required from you. Your changes are working as expected, well done! |
linalg promotion attempts to compute a constant upper bound for the allocated buffer size. Only when failed to compute an upperbound it fallbacks to the original subview size, which may be dynamic.
Adding a promotion option to use the original subview size by default, thus minimizing the allocation size.
Fixes #144268.